fix: preserve Unicode token chunk offsets - #646
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthrough
ChangesUnicode-safe token chunking
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant TokenChunker
participant TokenizerAdapter
participant SourceText
participant Chunk
TokenChunker->>TokenizerAdapter: encode_with_offsets(text)
TokenizerAdapter-->>TokenChunker: token IDs and character offsets
TokenChunker->>SourceText: select safe token-boundary spans
SourceText-->>TokenChunker: sliced chunk text
TokenChunker->>Chunk: store text, token count, indices, and overlap
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/chonkie/tokenizer.py (1)
676-680: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueCast
.idstolistfor consistency with the other adapters.
TiktokenAutoTokenizer.encode_with_offsetsandTransformersAutoTokenizer.encode_with_offsetsboth explicitly wrap their token IDs inlist(...).TokenizersAutoTokenizer.encode_with_offsets(Line 679) andTokieAutoTokenizer.encode_with_offsets(Line 697) returnencoding.idsdirectly. The consumer inchunker/token.py(_chunk_tokens) doeslist(tokens) == offset_tokens. If either binding's.idsis not a plainlist, this equality silently fails and the offset path silently falls back to decode-based chunking on every call, defeating the feature without any error.Cast explicitly for consistency and defensive robustness.
♻️ Proposed fix
def encode_with_offsets(self, text: str) -> tuple[list[int], list[tuple[int, int]]]: """Encode text and return character offsets for each token.""" encoding = self.tokenizer.encode(text, add_special_tokens=False) - return encoding.ids, encoding.offsets + return list(encoding.ids), list(encoding.offsets)Also applies to: 694-698
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/chonkie/tokenizer.py` around lines 676 - 680, Update encode_with_offsets in TokenizersAutoTokenizer and TokieAutoTokenizer to wrap encoding.ids in list(...) before returning it, while preserving the existing offsets unchanged and maintaining the declared tuple return shape.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/chonkie/chunker/token.py`:
- Around line 106-169: Restructure the tokenization flow so offset-aware
tokenizers are attempted first, using the tokens and offsets returned by
encode_with_offsets as the working data without re-encoding. Update chunk() and
_process_batch to avoid unconditional encode/encode_batch calls, while
preserving batched tokenization where offsets are unavailable or fail. Adjust
_chunk_tokens to consume the successful offset result directly and retain the
existing decode-based fallback.
In `@tests/chunkers/test_token_chunker.py`:
- Around line 278-280: Extend the tokenizer parametrization for
test_unicode_boundary_offsets to include the tokie fixture if
TokieAutoTokenizer.encode_with_offsets supports the expected offset API;
otherwise explicitly mark or document the fixture as unsupported for this test.
---
Nitpick comments:
In `@src/chonkie/tokenizer.py`:
- Around line 676-680: Update encode_with_offsets in TokenizersAutoTokenizer and
TokieAutoTokenizer to wrap encoding.ids in list(...) before returning it, while
preserving the existing offsets unchanged and maintaining the declared tuple
return shape.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 2b54765b-c8ef-4189-b493-141bca4f9e49
📒 Files selected for processing (3)
src/chonkie/chunker/token.pysrc/chonkie/tokenizer.pytests/chunkers/test_token_chunker.py
1 participant
Summary
Root cause
TokenChunker decoded each token group independently and accumulated decoded-string lengths. A group beginning in the middle of a UTF-8 character could decode as empty, which lost text and desynchronized later indices.
Validation
Summary by CodeRabbit